"""
A general implementation of particle swarm optimization (PSO).

PSO operates on a population of particles that fly through a solution space, in
search of the optimal position. Each particle keeps track of its personal best
position and velocity, which it uses along with the global best position to
calculate a new position and velocity. For more details see:
http://www.swarmintelligence.org/

Most functions in this module are higher-order functions that return another
function. To create a PSO optimizer, you must minimally call create_optimizer
with a fitness function that you must supply yourself.

Copyright (C) 2013 Jerrad Michael Genson

This program is free software: you can redistribute it and/or modify
it under the terms of the BSD 3-Clause License as published by
the Open Source Initiative.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
BSD 3-Clause License for more details.

You should have received a copy of the BSD 3-Clause License
along with this program.  If not, see
<https://betelgeuse-ml.googlecode.com/hg/LICENSE>

"""

# Python built-in modules
import random

# Third-party modules
import numpy as np

# Module header
__author__ = 'Jerrad Genson'
__contact__ = 'jerradgenson@gmail.com'
__copyright__ = 'Copyright 2013 Jerrad Michael Genson'
__license__ = 'BSD 3-Clause'

# Global constants
# Learning factors determine how quickly particles accelerate.
# The higher the learning factor, the faster the increase in velocity.
# Typically learning factor1 = learning factor2 = 2
DEFAULT_LEARNING_FACTOR1 = 2
DEFAULT_LEARNING_FACTOR2 = 2


def create_update_velocity(learning_factor1=DEFAULT_LEARNING_FACTOR1,
                           learning_factor2=DEFAULT_LEARNING_FACTOR2):
    """
    Create an update velocity function for use with create_optimizer.

    Args:
        learning_factor1: Value for the first learning factor (c1).
        learning_factor2: Value for the second learning factor (c2).

    Returns:
        An update velocity function.

    """

    def update_velocity(particle, global_best):
        """
        Calculates a new velocity for a given particle.

        Args:
            particle: The particle to calculate a new velocity for.
            global_best: The particle with the global best fitness (position).

        Returns:
            A new particle velocity.

        """

        new_velocity = (particle.velocity + learning_factor1 * random.random() *
                                            (particle.position - particle) +
                        learning_factor2 *
                        random.random() *
                        (global_best - particle))

        return new_velocity

    return update_velocity


def create_update_position():
    """
    Create an update particle position function for use with create_optimizer.

    Returns:
       An update position function - see documentation for update_position.

     """

    def update_position(particle):
        """
        Calculate a new position for a given particle.

        Args:
            particle The particle to calculate a new position for.

        Returns:
            The new particle position.

        """

        new_position = particle + particle.velocity
        return new_position

    return update_position


def create_optimizer(fitness, update_velocity=create_update_velocity(),
                     update_position=create_update_position(),
                     bounds=None):
    """
    Create a particle swarm optimizer function.

    Args:
        fitness: A function that accepts a particle as input and
                 returns its fitness as a real value between 0 and 1.

    Keywords:
        update_velocity: A function that accepts as inputs a particle and
                         the global best particle and returns a new
                         velocity for the given particle.
        update_position: A function that accepts a particle as its only
                         input and returns a new particle position.
        bounds: A tuple of (low, high) real values that set the bounds
                for a particle's velocity.

    Returns:
        A particle swarm optimization function.

    """

    def particle_swarm_optimizer(particles):
        """
        An optimizer function using particle swarm optimization.

        Args:
            particles: A list of Particle objects.

        Returns:
            A list of Particles after one cycle of PSO.

        """

        # Remember global best particle.
        global_best_particle = None
        # Calculate fitness of each particle.
        for particle in particles:
            particle_fitness = fitness(particle)
            if particle_fitness > particle.personal_best_fitness:
                # Current fitness better than particle's personal best.
                particle.personal_best_fitness = particle_fitness
                particle.position = particle.copy()

            if (global_best_particle is None or
                        particle.personal_best_fitness > global_best_particle.personal_best_fitness):
                # Particle's personal best fitness is better than global best.
                global_best_particle = particle.position

        new_particles = []
        # Calculate new velocity and position for each particle.
        for particle in particles:
            new_velocity = update_velocity(particle, global_best_particle)
            if bounds:
                # Make sure velocities are within given bounds.
                for index, element in enumerate(new_velocity):
                    if element < bounds[0]:
                        new_velocity[index] = bounds[0]

                    elif element > bounds[1]:
                        new_velocity[index] = bounds[1]

            particle.velocity = new_velocity
            new_particle = update_position(particle)
            new_particles.append(new_particle)

        return new_particles

    return particle_swarm_optimizer


class Particle(np.ndarray):
    """
    A Numpy Array wrapper with added attributes to be compatible with PSO.

    """

    def __new__(subtype, shape, dtype=float, buffer=None, offset=0,
                strides=None, order=None, info=None):
        # ndarray uses the __new__ constructor instead of the typical __init__,
        # so this is what we must override. 
        # http://docs.scipy.org/doc/numpy/user/basics.subclassing.html

        # Create the ndarray instance as usual.
        obj = np.ndarray.__new__(subtype, shape, dtype, buffer, offset,
                                 strides, order)

        # Append custom Particle attributes to the instance.
        obj.personal_best_fitness = 0
        if dtype:
            obj.position = np.zeros(shape, dtype=dtype)
            obj.velocity = np.zeros(shape, dtype=dtype)

        else:
            obj.position = np.array(shape)
            obj.velocity = np.zeros([len(shape)],
                                    dtype=obj.position.dtype)

        return obj

    def __array_finalize__(self, obj):
        # We have to override this method when subclassing ndarray to support
        # view casting and new-from-template.
        # http://docs.scipy.org/doc/numpy/user/basics.subclassing.html

        if obj is None:
            return

        # Copy custom Particle attributes into the new instance.
        self.personal_best_fitness = getattr(obj, 'personal_best_fitness')
        self.position = getattr(obj, 'position').copy()
        self.velocity = getattr(obj, 'velocity').copy()
